package org.dreamwork.tools.sdq;

import org.dreamwork.tools.sdq.conf.ServerConfig;
import org.dreamwork.tools.sdq.model.Endpoint;
import org.dreamwork.util.IDisposable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;

public class LoadBalancer implements IDisposable, Runnable {
    private final AtomicLong counter = new AtomicLong (0);
    private final Set<Endpoint> endpoints = new HashSet<> ();

    private final Object LOCKER = new byte[0];
    private final Thread thread = new Thread (this);

    private final long timeout, interval;
    private final Logger logger = LoggerFactory.getLogger (LoadBalancer.class);

    private transient int totalWeight = 0;
    private transient boolean running = true;

    public LoadBalancer (ServerConfig config) {
        this.timeout  = config.getHealthCheckTimeout ();
        this.interval = config.getHealthCheckInterval ();
    }

    public synchronized void register (Endpoint endpoint) {
        if (!endpoints.contains (endpoint)) {
            endpoints.add (endpoint);
            totalWeight += endpoint.weight;
            synchronized (LOCKER) {
                LOCKER.notifyAll ();
            }
        }
    }

    public synchronized void remove (Endpoint endpoint) {
        if (endpoints.contains (endpoint)) {
            endpoints.remove (endpoint);
            totalWeight -= endpoint.weight;
        }
    }

    public Endpoint getNextEndpoint () {
        int index = (int) (counter.incrementAndGet () % totalWeight);
        for (Endpoint ep : endpoints) {
            index -= ep.weight;
            if (index < 0) {
                return ep;
            }
        }
        return null;
    }

    public void start () {
        thread.start ();
    }

    @Override
    public void dispose () {
        running = false;
        thread.interrupt ();
    }

    @Override
    public void run () {
        while (running && !thread.isInterrupted ()) {
            while (endpoints.isEmpty ()) {
                synchronized (LOCKER) {
                    try {
                        LOCKER.wait ();
                    } catch (InterruptedException e) {
                        throw new RuntimeException (e);
                    }
                }
            }

            Set<Endpoint> toRemove = new HashSet<> ();
            long now = System.currentTimeMillis ();
            for (Endpoint ep : endpoints) {
                if (now - ep.lastHop > timeout) {
                    toRemove.add (ep);
                }
            }
            if (!toRemove.isEmpty ()) {
                for (Endpoint ep : toRemove) {
                    if (now - ep.lastHop > timeout) {
                        endpoints.remove (ep);
                    }
                }
            }

            try {
                Thread.sleep (interval);
            } catch (InterruptedException ex) {
                //
            }
        }
        logger.info ("LoadBalancer shutdown!");
    }
}